package com.pnnd;

import com.pnnd.utils.ArrayUtils;
import com.pnnd.utils.IntTuple;

import java.lang.reflect.Array;
import java.util.Iterator;

public class NdArrayFlatIter<E> implements Iterable<E> {

    private Object ndArray;
    private IntTuple tuple;

    public NdArrayFlatIter(Object ndArray){
        this.ndArray = ndArray;

        this.tuple = ArrayUtils.getShape(ndArray);
    }

    @Override
    public Iterator<E> iterator() {
        return new NdArrayIterator(this);
    }

    public class NdArrayIterator implements Iterator<E>{

        private int[] dims;
        private int[] indices;
        private NdArrayFlatIter iter;

        public NdArrayIterator(NdArrayFlatIter<E> iter){
            this.iter = iter;
            this.dims = iter.tuple.toArray();

            this.indices = IntTuple.fullNum(iter.tuple.size(),0).toArray();
        }


        @Override
        public boolean hasNext() {

            for(int i=0;i<this.indices.length;i++){
                if(this.indices[i] <= this.dims[i] - 1){
                    return true;
                }else{
                    return false;
                }
            }

            return false;
        }

        @Override
        public E next() {


            try {
                Object arrObj = this.iter.ndArray;

                for(int i=0;i<indices.length;i++){
                    arrObj = Array.get(arrObj,this.indices[i]);
                    if(arrObj == null){
                        return null;
                    }
                }

                return (E)arrObj;
            }finally {
                this.indexIncreasing(this.indices.length);
            }
        }

        private void indexIncreasing(int dim){

            this.indices[dim - 1] += 1;
            if(this.indices[dim - 1] == this.dims[dim - 1] && dim > 1){
                this.indices[dim - 1] = 0;
                this.indexIncreasing(dim - 1);
            }
        }
    }

    public static void main(String[] args){
        int[][] arr = new int[2][3];

        arr[0] = new int[]{1,2,3};
        arr[1] = new int[]{4,5,6};

        for(int o:new NdArrayFlatIter<Integer>(arr)){
            System.out.println(o);
        }
    }
}
